import requests
import json
import base64
import os
from io import BytesIO
from PIL import Image
import time

class TextToImageGenerator:
    def __init__(self):
        """
        初始化ModelScope API文本生图生成器
        使用Qwen-Image模型通过API调用实现更快的生成速度
        """
        self.base_url = 'https://api-inference.modelscope.cn/'
        self.api_key = os.environ.get("MODELSCOPE_API_TOKEN")
        self.model_id = "Qwen/Qwen-Image"
        
        self.common_headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        
        print("ModelScope API文本生图模型初始化成功")
    
    def generate_image(self, prompt, size="512x512", num_images=1, style="photographic"):
        """
        根据文本生成图像
        Args:
            prompt (str): 文本描述，支持中英文
            size (str): 图像尺寸，支持"512x512", "768x512", "1024x1024"等
            num_images (int): 生成图像数量(1-4)
            style (str): 图像风格，可选"photographic", "digital-art", "anime", "oil-painting"
        Returns:
            dict: 包含生成结果的字典
        """
        try:
            # 限制生成数量
            num_images = min(max(1, num_images), 4)
            
            # 解析尺寸
            width, height = map(int, size.split('x'))
            
            # 构建请求数据
            request_data = {
                "model": self.model_id,
                "prompt": prompt,
                "n": num_images,
                "size": f"{width}x{height}",
                "style": style,
                "response_format": "b64_json"  # 返回base64格式
            }
            
            # 发送异步生成请求
            response = requests.post(
                f"{self.base_url}v1/images/generations",
                headers={**self.common_headers, "X-ModelScope-Async-Mode": "true"},
                data=json.dumps(request_data, ensure_ascii=False).encode('utf-8')
            )
            
            response.raise_for_status()
            task_data = response.json()
            
            if "task_id" not in task_data:
                raise Exception(f"API响应异常: {task_data}")
            
            task_id = task_data["task_id"]
            
            # 轮询任务状态
            images = self._poll_task_status(task_id)
            
            return {
                "success": True,
                "images": images,
                "prompt": prompt,
                "size": size,
                "style": style,
                "model": "Qwen-Image"
            }
            
        except requests.exceptions.RequestException as e:
            return {
                "success": False,
                "error": f"网络请求错误: {str(e)}",
                "images": []
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "images": []
            }
    
    def _poll_task_status(self, task_id, max_wait=300):
        """
        轮询任务状态直到完成或超时
        
        Args:
            task_id: 任务ID
            max_wait: 最大等待时间(秒)
            
        Returns:
            list: 生成的图像列表
        """
        start_time = time.time()
        
        while time.time() - start_time < max_wait:
            try:
                result = requests.get(
                    f"{self.base_url}v1/tasks/{task_id}",
                    headers={**self.common_headers, "X-ModelScope-Task-Type": "image_generation"},
                )
                
                result.raise_for_status()
                data = result.json()
                
                if data["task_status"] == "SUCCEED":
                    images = []
                    
                    # 处理生成的图片
                    for img_url in data.get("output_images", []):
                        try:
                            # 下载图片
                            img_response = requests.get(img_url)
                            img_response.raise_for_status()
                            
                            # 转换为base64
                            img = Image.open(BytesIO(img_response.content))
                            buffered = BytesIO()
                            img.save(buffered, format="PNG")
                            img_b64 = base64.b64encode(buffered.getvalue()).decode()
                            
                            images.append({
                                "url": f"data:image/png;base64,{img_b64}",
                                "size": "512x512"  # API返回固定尺寸
                            })
                            
                        except Exception as e:
                            print(f"处理图片时出错: {e}")
                            continue
                    
                    return images
                    
                elif data["task_status"] == "FAILED":
                    error_msg = data.get("message", "未知错误")
                    raise Exception(f"生成失败: {error_msg}")
                
                # 等待5秒后重试
                time.sleep(5)
                
            except requests.exceptions.RequestException as e:
                raise Exception(f"查询任务状态时出错: {str(e)}")
        
        raise Exception("任务超时，请稍后重试")
    
    def get_available_sizes(self):
        """获取支持的图像尺寸"""
        return [
            "512x512",
            "768x512",
            "512x768", 
            "768x768",
            "1024x512",
            "512x1024",
            "1024x1024"
        ]
    
    def get_available_styles(self):
        """获取支持的图像风格"""
        return [
            "photographic",
            "digital-art",
            "anime",
            "oil-painting",
            "sketch",
            "3d-model"
        ]
    
    def check_api_status(self):
        """检查API状态"""
        try:
            response = requests.get(
                f"{self.base_url}v1/models",
                headers=self.common_headers
            )
            
            if response.status_code == 200:
                models = response.json()
                qwen_available = any(
                    model.get("id") == self.model_id 
                    for model in models.get("data", [])
                )
                
                return {
                    "available": qwen_available,
                    "message": "Qwen-Image模型可用" if qwen_available else "Qwen-Image模型不可用"
                }
            else:
                return {
                    "available": False,
                    "message": f"API检查失败: {response.status_code}"
                }
                
        except Exception as e:
            return {
                "available": False,
                "message": f"连接失败: {str(e)}"
            }

# 创建全局实例
image_generator = TextToImageGenerator()
